MODULE optimize_primitives
  !---------------------------------------------------------------------
  ! Completeness-optimization procedures used by Kruununhaka
  !---------------------------------------------------------------------
  ! Copyright (C) 2004-2007 Pekka Manninen, University of Helsinki
  !                    2010 Jussi Lehtola,  University of Helsinki
  !
  ! This program is distributed under the terms of the GNU General
  ! Public License as published by the Free Software Foundation;
  ! either version 2 of the License, or (at your option) any later version.
  !
  ! This program is distributed in the hope that it will be useful,
  ! but WITHOUT ANY WARRANTY; without even the implied warranty of
  ! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. In addition, the 
  ! author must be properly cited whenever the program or some parts 
  ! originated on it are used. 
  !
  !---------------------------------------------------------------------
  USE definitions
  USE error
  USE random
  USE completeness_profile, ONLY : comp_prof
  USE measure
  USE sorting

  IMPLICIT NONE
  PRIVATE
  PUBLIC :: opt_prim, opt_prim_bf

CONTAINS
!---------------------------------------------------------------------
  SUBROUTINE opt_prim(l, optmin, optmax, X, C, usesto, measure)
    ! Optimize the primitives
    IMPLICIT NONE
    ! Value of angular momentum
    INTEGER, INTENT(IN) :: l
    ! Exponent range to optimize in
    REAL(KIND=prec), INTENT(IN) :: optmin, optmax
    ! Exponents to be optimized
    REAL(KIND=prec), DIMENSION(:), INTENT(OUT) :: X
    ! Contraction coefficients
    REAL(KIND=prec), DIMENSION(:,:), INTENT(IN) :: C
    ! Use STOs to probe?
    LOGICAL, INTENT(IN) :: usesto
    ! Which measure to use
    INTEGER, INTENT(IN) :: measure
    ! Probing exponents
    REAL(KIND=prec), DIMENSION(optpoints) :: alphas
    ! Evolution populations
    REAL(KIND=prec), DIMENSION(SIZE(X),pop_dim*SIZE(X)) :: popul
    ! Values of cost function for the populations
    REAL(KIND=prec), DIMENSION(pop_dim*SIZE(X)) :: cost
    ! Time variables
    INTEGER :: time1, time2, cpucount
    ! Amount of exponents
    INTEGER :: nx
    ! Location of minimum cost
    INTEGER, DIMENSION(1) :: minind
    nx=SIZE(X)

    ! Fill probing exponents
    CALL get_alphas(optmin,optmax,alphas)
    ! Print values
    WRITE(*,'(/,A)') 'Values for initial set:'
    CALL f_everything(l, X, C, alphas, usesto)
    WRITE(*,'(A,I3,A)') 'Optimizing a set of ', nx, &
         ' functions. This may take time.'

    CALL system_clock(time1,cpucount)
    ! Make a population by mutating the optimal exponents by scaling them
    CALL oned_init(l, measure, X, C, popul, cost, alphas, usesto)
    ! Run the optimization on the mutated population
    CALL oned_opt(l, measure, C, popul, cost, alphas, usesto)
    CALL system_clock(time2,cpucount)

    WRITE (*,'(A,F5.1,A,/)') 'Optimization took',&
         (time2-time1)*1.0/cpucount,' seconds.'
    ! Get minimum cost value
    minind=MINLOC(cost)
    ! Get optimal set
    X(1:nx)=popul(1:nx,minind(1))
    ! Sort it in descending order
    CALL sort_desc(X)

    ! Print values
    WRITE(*,'(A)') 'Values for optimized set:'
    CALL f_everything(l, X, C, alphas, usesto)

  END SUBROUTINE opt_prim
!---------------------------------------------------------------------
  SUBROUTINE oned_init(l, measure, opt_set, C, pop, cost, alphas, usesto)
    IMPLICIT NONE
    ! Angular momentum type
    INTEGER, INTENT(IN) :: l 
    ! Measure to use in optimization
    INTEGER, INTENT(IN) :: measure
    ! Current guess of optimal exponents
    REAL(KIND=prec), DIMENSION(:), INTENT(IN) :: opt_set
    ! Contraction coefficients
    REAL(KIND=prec), DIMENSION(:,:), INTENT(IN) :: C
    ! Population  
    REAL(KIND=prec), DIMENSION(:,:), INTENT(OUT) :: pop
    ! Values of cost function for population
    REAL(KIND=prec), DIMENSION(:), INTENT(OUT) :: cost
    ! Scanning exponents
    REAL(KIND=prec), DIMENSION(:), INTENT(IN) :: alphas
    ! Use STO to probe?
    LOGICAL, INTENT(IN) :: usesto
    ! Number of exponents and population size
    INTEGER :: nx, pop_size
    ! Loop indices
    INTEGER :: i, j

    ! Number of exponents to optimize
    nx = SIZE(opt_set)
    ! Population size
    pop_size = SIZE(pop,2)
    IF(SIZE(cost) /= pop_size) STOP 'Cost vector of wrong size'

    ! Form new population
    SELECT CASE(measure)
    CASE(1) ! Use maximum deviance
       DO i = 1, pop_size
          cost(i)=1.0
          ! Try to get rid of pathogenic cases
          DO WHILE(cost(i)>max_path)
             ! Mutate the exponents
             pop(:,i) = evol_low + (evol_up-evol_low)*rnd_vec(nx)
             DO j = 1, nx
                pop(j,i) = pop(j,i)*opt_set(j)
             END DO

             cost(i) = f_maxdev(l,pop(:,i),C,alphas,usesto)
          END DO
       END DO

    CASE(2) ! Use squared deviance
       DO i = 1, pop_size
          cost(i)=1.0
          ! Try to get rid of pathogenic cases
          DO WHILE(cost(i)>max_path)
             ! Mutate the exponents
             pop(:,i) = evol_low + (evol_up-evol_low)*rnd_vec(nx)
             DO j = 1, nx
                pop(j,i) = pop(j,i)*opt_set(j)
             END DO

             cost(i) = f_lsq(l,pop(:,i),C,alphas,usesto)
          END DO
       END DO

    CASE(3) ! Use integral
       DO i = 1, pop_size
          cost(i)=1.0
          ! Try to get rid of pathogenic cases
          DO WHILE(cost(i)>max_path)
             ! Mutate the exponents
             pop(:,i) = evol_low + (evol_up-evol_low)*rnd_vec(nx)
             DO j = 1, nx
                pop(j,i) = pop(j,i)*opt_set(j)
             END DO
             cost(i) = f_sum(l,pop(:,i),C,alphas,usesto)
          END DO
       END DO
    END SELECT
  END SUBROUTINE oned_init
!---------------------------------------------------------------------
  SUBROUTINE oned_opt(l, measure, C, pop, cost, alphas, usesto)
    ! optimization routine for vector arrays
    IMPLICIT NONE
    ! Angular momentum
    INTEGER, INTENT(IN) :: l
    ! Measure to minimize
    INTEGER, INTENT(IN) :: measure
    ! Contraction coefficients
    REAL(KIND=prec), DIMENSION(:,:), INTENT(IN) :: C
    ! The population
    REAL(KIND=prec), DIMENSION(:,:), INTENT(INOUT) :: pop
    ! Cost functions for the populations
    REAL(KIND=prec), DIMENSION(:), INTENT(INOUT) :: cost
    ! Scanning exponents
    REAL(KIND=prec), DIMENSION(:), INTENT(IN) :: alphas
    ! Use STO to probe?
    LOGICAL, INTENT(IN) :: usesto
    ! Number of exponents and population size
    INTEGER :: nx, pop_size
    ! Helper indices
    INTEGER ::  allocstat, iter, i, j, idx1, idx2, idx3
    ! New population
    REAL(KIND=prec), DIMENSION(:,:), ALLOCATABLE :: new_pop
    ! Minimum cost functions for old iterations
    REAL(KIND=prec), DIMENSION(:), ALLOCATABLE :: oldit
    ! Trial exponents
    REAL(KIND=prec), DIMENSION(SIZE(pop,1)) :: trial
    ! Values of cost function
    REAL(KIND=prec) :: score, itthr
    ! RND variable for determining which exponents to pair
    REAL(KIND=prec), DIMENSION(SIZE(pop,1)) :: p
    ! Has the evolution converged?
    LOGICAL :: converged

    nx = SIZE(pop,1)
    pop_size = SIZE(pop,2)
    IF(SIZE(cost) /= pop_size) STOP 'Vector cost is of wrong size'
    ! Allocate new population
    ALLOCATE(new_pop(nx,pop_size), STAT=allocstat)
    IF(allocstat /= 0) STOP 'Error in memory allocation (oned_opt)'

    ! Inilize old population
    new_pop = 0.0_prec
    score=0.0_prec
    ! Allocate memory for minimal costs of old iterations
    ALLOCATE(oldit((maxit/istep)+1), STAT=allocstat)
    IF(allocstat /= 0) STOP 'Error in memory allocation (oned_opt)'
    oldit = 0.0_prec
    ! Index of old iterations
    j = 2
    ! Current iteration number
    iter = 1
    converged = .FALSE.

    SELECT CASE(measure)
    CASE(1)
       WRITE(*,'(/,A10,A10,A15)') ' ', 'Iteration', 'Max deviation'
    CASE(2)
       WRITE(*,'(/,A10,A10,A15)') ' ', 'Iteration', 'RMS deviation'
    CASE(3)
       WRITE(*,'(/,A10,A10,A15)') ' ', 'Iteration', 'Integral'
    END SELECT

    DO WHILE (iter <= maxit .AND. .NOT. converged)
       ! Form new population
       DO i = 1, pop_size
          ! Get three random indices of populations to couple
          CALL rnd_triple(i,pop_size,idx1,idx2,idx3)

          ! Get random numbers to determine which exponents to couple
          p = rnd_vec(nx)
          WHERE (p<evol_r)
             trial = ABS(pop(:,idx3) + evol_c*(pop(:,idx1)-pop(:,idx2)))
          ELSEWHERE
             trial = ABS(pop(:,i))
          END WHERE

          ! Calculate cost
          SELECT CASE(measure)
          CASE(1)
             ! Use maximum deviance
             score = f_maxdev(l,trial,C,alphas,usesto)
          CASE(2)
             ! Use squared deviance
             score = f_lsq(l,trial,C,alphas,usesto)
          CASE(3)
             ! Use integral
             score = f_sum(l,trial,C,alphas,usesto)
          CASE DEFAULT
             CALL error_stop('Error in oned_opt - invalid measure!')
          END SELECT
          ! Is the new trial better than the one we are currently 
          ! iterating over?
          IF(score < cost(i)) THEN
             ! Yes, replace it with the trial
             new_pop(:,i) = trial
             cost(i) = score
             itthr = MINVAL(cost)
          ELSE
             ! No, let old specimen in new population
             new_pop(:,i) = pop(:,i)
          END IF
       END DO
       ! Change populations.
       pop = new_pop
       ! Printout if necessary
       IF(MOD(iter,istep) == 0) THEN
          WRITE(*,'(A10,I10,E15.6)') ' ', iter, MINVAL(cost)
          oldit(j) = MINVAL(cost)
          IF (ABS(oldit(j) - oldit(j-1)) < conv_thr) THEN
             converged = .TRUE.
          END IF
          j = j + 1
       END IF
       iter = iter + 1
    END DO

    IF (iter < maxit) THEN
       WRITE(*,'(A,I5,A)') 'Optimization converged after ',iter,' iterations!'
    ELSE
       WRITE(*,'(A)') 'Maximum number of iterations reached.'
    END IF
    DEALLOCATE(new_pop)
    DEALLOCATE(oldit)  
  END SUBROUTINE oned_opt
!---------------------------------------------------------------------
  FUNCTION opt_prim_bf(l, optmin, optmax, X, C, usesto, measure, verbose) RESULT(curiter)
    ! Optimize the primitives by brute force
    IMPLICIT NONE
    ! Value of angular momentum
    INTEGER, INTENT(IN) :: l
    ! Exponent range to optimize in
    REAL(KIND=prec), INTENT(IN) :: optmin, optmax
    ! Exponents to optimize
    REAL(KIND=prec), DIMENSION(:), INTENT(INOUT) :: X
    ! Contraction coefficients
    REAL(KIND=prec), DIMENSION(:,:), INTENT(IN) :: C
    ! Use STOs to probe?
    LOGICAL, INTENT(IN) :: usesto
    ! Which measure to use
    INTEGER, INTENT(IN) :: measure
    ! Verbose operation?
    LOGICAL, INTENT(IN) :: verbose

    ! Probing exponents
    REAL(KIND=prec), DIMENSION(optpoints) :: alphas
    
    ! Step sizes
    REAL(KIND=prec), DIMENSION(SIZE(X)) :: stepsize
    ! RMS step
    REAL(KIND=prec) :: rmsstep
    ! When to break off due to small gradient
    REAL(KIND=prec), PARAMETER :: eps=EPSILON(stepsize)
    ! When to break off due to small total change of cost
    REAL(KIND=prec), PARAMETER :: tol=1e-30
    ! Current iteration
    INTEGER :: iter
    
    ! Time variables
    INTEGER :: time1, time2, cpucount
    ! Amount of exponents
    INTEGER :: nx
    
    ! Loop indices
    INTEGER :: ix, io, j
    ! Loop order exponents
    INTEGER, DIMENSION(SIZE(X)) :: loopord
    ! Logarithms of exponents
    REAL(KIND=prec), DIMENSION(SIZE(X)) :: logX
    ! Trial exponents
    REAL(KIND=prec), DIMENSION(SIZE(X),-1:1) :: Xtrial
    
    ! Costs of current and last iteration over all exponents
    REAL(KIND=prec) :: curiter, lastiter
    ! Current cost
    REAL(KIND=prec) :: curcost=0.0_prec
    ! Costs of varying current exponent
    REAL(KIND=prec), DIMENSION(-1:1) :: trialcost

    ! Position and value of minimum
    INTEGER :: minpos
    REAL(KIND=prec) :: minval

    curiter=0.0_prec
    nx=SIZE(X)

    ! Fill probing exponents
    CALL get_alphas(optmin,optmax,alphas)
    ! Print values
    IF(verbose) THEN
       WRITE(*,'(/,A)') 'Values for initial set:'
       CALL f_everything(l, X, C, alphas, usesto)
       WRITE(*,'(A,I3,A)') 'Optimizing a set of ', nx, &
            ' functions. This may take time.'
    END IF

    CALL system_clock(time1,cpucount)

    ! Fill looping order. It is wise to move the outer exponents
    ! first, since the middle exponent range is crowded.
    loopord=0

    j=1
    DO ix=1,nx/2
       loopord(j)=ix
       loopord(j+1)=nx+1-ix
       j=j+2
    END DO
    ! If nx is odd, we need to add in the last exponent (i.e. middle one) as well.
    IF(MOD(nx,2)==1) THEN
       loopord(nx)=nx/2+1
    END IF

    ! Calculate logarithms of exponents
    logX=LOG(X)
    ! Initial step sizes
    stepsize=1.0_prec

    SELECT CASE(measure)
    CASE(1)
       IF(verbose) WRITE(*,'(/,A10,A10,A15,A15)') ' ', 'Iteration', 'Max deviation','RMS step size'
       curiter=f_maxdev(l,X,C,alphas,usesto)
       curcost=curiter
    CASE(2)
       IF(verbose) WRITE(*,'(/,A10,A10,A15,A15)') ' ', 'Iteration', 'RMS deviation','RMS step size'
       curiter=f_lsq(l,X,C,alphas,usesto)
       curcost=curiter
    CASE(3)
       IF(verbose) WRITE(*,'(/,A10,A10,A15,A15)') ' ', 'Iteration', 'Integral','RMS step'
       curiter=f_sum(l,X,C,alphas,usesto)
       curcost=curiter
    END SELECT

    iter=1
    ! Optimize.
    DO WHILE(MAXVAL(stepsize)>eps)
       ! Increase step sizes
       stepsize=2.0_prec*stepsize

       ! Store old value of cost function
       lastiter=curiter

       ! Loop over exponents
       DO io=1,nx
          ! Current exponent is
          ix=loopord(io)

          ! Current exponents
          Xtrial(:,0)=X

          DO WHILE(stepsize(ix)>eps)
             ! Calculate trials
             DO j=-1,1,2
                Xtrial(:,j)=X
                Xtrial(ix,j)=EXP(logX(ix)+j*stepsize(ix))
             END DO

             ! and costs
             SELECT CASE(measure)
                CASE(1)
                   DO j=-1,1,2
                      trialcost(j)=f_maxdev(l,Xtrial(:,j),C,alphas,usesto)
                   END DO
                   trialcost(0)=curcost
                CASE(2)
                   DO j=-1,1,2
                      trialcost(j)=f_lsq(l,Xtrial(:,j),C,alphas,usesto)
                   END DO
                   trialcost(0)=curcost
                CASE(3)
                   DO j=-1,1,2
                      trialcost(j)=f_sum(l,Xtrial(:,j),C,alphas,usesto)
                   END DO
                   trialcost(0)=curcost
              END SELECT

              ! Find out minimum position
              minval=trialcost(0)
              minpos=0
              DO j=-1,1,2
                 IF(trialcost(j)<minval) THEN
                    minpos=j
                    minval=trialcost(j)
                 END IF
              END DO
              
              ! Does stepsize need adjusting?
              IF(minpos==0) THEN
                 stepsize(ix)=stepsize(ix)/2.0_prec
              ELSE
                 ! No, break while loop
                 EXIT
              END IF
           END DO

           ! If all went well, we have a lower value
           IF(trialcost(-1)<curcost) THEN
              curcost=trialcost(-1)
              X=Xtrial(:,-1)
              logX(ix)=logX(ix)-stepsize(ix)
           ELSE IF(trialcost(1)<curcost) THEN
              ! Lower value is on the right
              curcost=trialcost(1)
              X=Xtrial(:,1)
              logX(ix)=logX(ix)+stepsize(ix)
!           ELSE
!              WRITE (*,*) 'Exponent ',ix,' is stagnated.'
           END IF
        END DO
        
     ! Store current value of cost
     curiter=curcost
     
     ! Printout if necessary
     IF(MOD(iter,istep_even) == 1) THEN
        ! Compute rms step size
        rmsstep=0.0_prec
        DO j=1,nx
           rmsstep=rmsstep + stepsize(j)*stepsize(j)
        END DO
        rmsstep=SQRT(rmsstep)/nx

        IF(verbose) WRITE(*,'(A10,I10,E15.6,E15.6)') ' ', iter, curiter, rmsstep        
     END IF
     iter=iter+1
     
     ! Is the change in the measure significant?
     IF(ABS(curiter-lastiter)<tol) THEN
        EXIT
     END IF
  END DO
  
  CALL system_clock(time2,cpucount)
  
  ! Sort optimal set in descending order
  CALL sort_desc(X)
  
  ! Print values
  IF(verbose) THEN
     WRITE (*,'(A,F5.1,A,/)') 'Optimization took',&
          (time2-time1)*1.0/cpucount,' seconds.'
     WRITE(*,'(A)') 'Values for optimized set:'
     CALL f_everything(l, X, C, alphas, usesto)
  END IF

END FUNCTION opt_prim_bf
!---------------------------------------------------------------------
END MODULE optimize_primitives
